#include<iostream>
#include<vector>

using namespace std;

class Solution {
public:
	void sortColors(vector<int>& nums) {
		int zeroCount = 0, oneCount = 0, twoCount = 0;
		for (int i = 0; i < nums.size(); i++) {
			if (nums[i] == 0) {
				zeroCount++;
			}
			else if (nums[i] == 1) {
				oneCount++;
			}
			else {
				twoCount++;
			}
			if (i == nums.size() - 1) {
				for (int j = 0; j < zeroCount; j++) {
					nums[j] = 0;
				}
				for (int j = zeroCount; j < zeroCount + oneCount; j++) {
					nums[j] = 1;
				}
				for (int j = zeroCount + oneCount; j < nums.size(); j++)
					nums[j] = 2;
			}
		}
	}

	void sortColors(vector<int>& nums) {
		int size = nums.size();
		if (size < 2) return;
		
		//all in [0,p0)==0
		//all in [p0,i)==1
		//all in (p2,size-1]==2
		int p0 = 0, i = 0, p2 = size - 1;
		while (i <= p2) {
			if (nums[i] == 0) {
				swap(nums[i], nums[p0]);
				p0++;
				i++;
			}
			else if (nums[i] == 1) {
				i++;
			}
			else {
				swap(nums[i], nums[p2]);
				p2--;
			}
		}
	}
};

int main() {
	Solution s;
	vector<int> arr = { 2,1 };
	s.sortColors(arr);
	return 0;
}